package com.snailwu.redis.semaphore;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.net.NetUtil;
import cn.hutool.core.text.StrPool;
import cn.hutool.core.util.RandomUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.ReturnType;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @author WuQinglong
 */
@Slf4j
public class RedisSemaphore {

    private final String key;
    private final String permits;
    private final String timeout;
    private final StringRedisTemplate stringRedisTemplate;
    private final ScheduledThreadPoolExecutor renewThreadPool = new ScheduledThreadPoolExecutor(1);
    private final ConcurrentHashMap<String, ScheduledFuture<?>> scheduledFutureMap = new ConcurrentHashMap<>();

    private static final ThreadLocal<String> permitUidMap = new ThreadLocal<>();
    private static final String acquireLuaScript;
    private static final String renewLuaScript;
    private static final String IP = NetUtil.getLocalHostName();

    static {
        try {
            ClassPathResource scriptResource = new ClassPathResource("/semaphore-acquire.lua");
            acquireLuaScript = IoUtil.readUtf8(scriptResource.getInputStream());
        } catch (IOException e) {
            throw new RuntimeException("解析信号量lua脚本异常", e);
        }

        try {
            ClassPathResource scriptResource = new ClassPathResource("/semaphore-renew.lua");
            renewLuaScript = IoUtil.readUtf8(scriptResource.getInputStream());
        } catch (IOException e) {
            throw new RuntimeException("解析信号量lua脚本异常", e);
        }
    }

    public RedisSemaphore(String key, int permits, long timeout, StringRedisTemplate stringRedisTemplate) {
        Assert.isTrue(timeout > 1);
        this.key = key;
        this.permits = String.valueOf(permits);
        this.timeout = String.valueOf(timeout);
        this.stringRedisTemplate = stringRedisTemplate;
        renewThreadPool.setCorePoolSize(permits);
    }

    /**
     * 尝试获取令牌
     *
     * @return true:获取成功；false:获取失败
     */
    public boolean tryAcquire() {
        String threadUid = threadUid();
        if (tryAcquire(threadUid)) {
            permitUidMap.set(threadUid);
            return true;
        }
        return false;
    }

    /**
     * 尝试获取令牌
     * @param permitUid 持有令牌者的唯一标识
     * @return true:获取成功；false:获取失败
     */
    public boolean tryAcquire(String permitUid) {
        try (RedisConnection connection = stringRedisTemplate.getRequiredConnectionFactory().getConnection()) {
            List<String> keysList = Collections.singletonList(key);
            List<String> argsList = Arrays.asList(permitUid, permits, System.currentTimeMillis() + "", timeout);
            byte[][] keysAndArgs = getKeyArgs(keysList, argsList);
            Boolean ret = connection.eval(acquireLuaScript.getBytes(StandardCharsets.UTF_8), ReturnType.BOOLEAN,
                    keysList.size(), keysAndArgs);
            if (Boolean.TRUE.equals(ret)) {
                log.info("获得信号量。key:{} permitUid:{}", key, permitUid);
                refreshExpireTime(permitUid);
                return true;
            }
        } catch (Exception e) {
            log.error("获取信号量令牌异常", e);
        }
        return false;
    }

    public void release() {
        // 分布式环境下获取这个线程唯一的一个标识
        String permitId = permitUidMap.get();
        permitUidMap.remove();

        release(permitId);
    }

    public void release(String permitId) {
        Optional.ofNullable(scheduledFutureMap.remove(permitId))
                .ifPresent(future -> future.cancel(true));
        stringRedisTemplate.opsForZSet().remove(key, permitId);
        log.info("释放信号量。key:{} permitId:{}", key, permitId);
    }

    /**
     * 定时续期
     */
    private void refreshExpireTime(String permitId) {
        long delay = (long) (Integer.parseInt(timeout) * 0.8);
        ScheduledFuture<?> scheduledFuture = renewThreadPool.scheduleWithFixedDelay(() -> {
            try (RedisConnection connection = stringRedisTemplate.getRequiredConnectionFactory().getConnection()) {
                List<String> keysList = Collections.singletonList(key);
                List<String> argsList = Arrays.asList(permitId, timeout);
                byte[][] keysAndArgs = getKeyArgs(keysList, argsList);
                Boolean ret = connection.eval(renewLuaScript.getBytes(StandardCharsets.UTF_8), ReturnType.BOOLEAN,
                        keysList.size(), keysAndArgs);
                if (Boolean.FALSE.equals(ret)) {
                    log.warn("续期信号量。{} {} 不存在，取消续期，移除续期线程", key, permitId);
                    Optional.ofNullable(scheduledFutureMap.remove(permitId))
                            .ifPresent(future -> future.cancel(true));
                } else {
                    log.info("续期信号量。key:{} permitId:{}", key, permitId);
                }
            } catch (Exception e) {
                log.error("定时续期信号量异常", e);
            }
        }, delay, delay, TimeUnit.MILLISECONDS);

        scheduledFutureMap.put(permitId, scheduledFuture);
    }

    /**
     * 获取分布式下线程的唯一id
     * {podIp}:{threadId}:{随机数}
     */
    private String threadUid() {
        long threadId = Thread.currentThread().getId();
        return IP + StrPool.COLON + threadId + StrPool.COLON + RandomUtil.randomString(5);
    }

    private static byte[][] getKeyArgs(List<String> keyList, List<String> argList) {
        byte[][] keyArgs = new byte[keyList.size() + argList.size()][];
        int i = 0;
        for (String key : keyList) {
            keyArgs[i++] = key.getBytes(StandardCharsets.UTF_8);
        }
        for (String arg : argList) {
            keyArgs[i++] = arg.getBytes(StandardCharsets.UTF_8);
        }
        return keyArgs;
    }

}
